"""
神经网络层相关
"""
import os.path

import numpy as np
import dezero
import dezero.functions as F
from dezero.core import Parameter
import weakref
from dezero import cuda

class Layer:
    """"神经网络层 基类"""
    def __init__(self):
        self._params = set() #_params保存管理所有参数

    def __setattr__(self, name, value):
        """实现在添加实例变量时，可以添加parameter类型 也可以添加layer类型嵌套"""
        if isinstance(value, (Parameter,Layer)): #判断是否是parameter或Layer类型,Layer可以嵌套
            self._params.add(name)
        super().__setattr__(name,value)

    def __call__(self, *inputs):
        outputs = self.forward(*inputs)
        if not isinstance(outputs,tuple):
            outputs = (outputs,)
        self.inputs = [weakref.ref(x) for x in inputs]
        self.outputs = [weakref.ref(y) for y in outputs]
        return outputs if len(outputs)>1 else outputs[0] #如果输出只有一个值则直接返回该值

    def forward(self,x):
        """前向传播的方法"""
        raise NotImplementedError()
    def params(self):
        """返回一个生成器"""
        for name in self._params:
            obj = self.__dict__[name]
            #所有变量都以字典形式存储在实例变量__dict__中，可以通过它直接取出parameter或Layer实例
            if isinstance(obj,Layer):
                yield from obj.params() #从嵌套的Layer递归取出函数
            else:
                yield obj
            #yield是暂停处理并返回值，此处和for语句可以按顺序取出参数。使用yield from可以使用一个生成器创建另一个新的生成器
    def cleargrads(self):
        """重置所有参数导数"""
        for param in self.params():
            param.cleargrad()

    def _flatten_params(self, params_dict, parent_key=""):
        """在层结构中将Parameter作为一个扁平的、非嵌套的字典取出"""
        for name in self._params:
            obj = self.__dict__[name]
            #若存在上一层则使用parent/name取名
            key = parent_key + '/' + name if parent_key else name

            if isinstance(obj, Layer):
                obj._flatten_params(params_dict, key)
            else:
                params_dict[key] = obj

    def save_weights(self, path):
        """保存参数 path：保存路径"""
        #self.to_cpu()
        #创建空字典给到函数，将参数递归填入,然后再将所有key和data放到array_dict中
        params_dict = {}
        self._flatten_params(params_dict)
        array_dict = {key: param.data \
                      for key,param in params_dict.items() \
                      if param is not None}
        try:
            np.savez_compressed(path, **array_dict)
        except (Exception, KeyboardInterrupt) as e:
            """若发生错误或者用户发起键盘中断，则停止保存并删除不完整的文件"""
            if os.path.exists(path):
                os.remove(path)
            raise

    def load_weights(self, path):
        """加载参数 path：加载路径"""
        npz = np.load(path)
        params_dict = {}
        self._flatten_params(params_dict)
        #将参数加载到一个已知的模型中
        for key, param in params_dict.items():
            param.data = npz[key]

    def to_cpu(self):
        """cupy->numpy"""
        for param in self.params():
            param.to_cpu()
    def to_gpu(self):
        """numpy->cupy"""
        for param in self.params():
            param.to_gpu()




class Linear(Layer):
    def __init__(self,out_size,nobias=False,dtype=np.float32, in_size=None):
        """
        参数：设置输入输出大小，nobias为True则省略偏置, 可以不指定输入大小会自动判断
        """
        super().__init__()
        self.in_size = in_size
        self.out_size = out_size
        self.dtype = dtype

        self.W = Parameter(None, name='W')
        if self.in_size is not None: #如果没有指定in_size,则延后处理
            self._init_W()

        if nobias:
            self.b = None
        else:
            self.b = Parameter(np.zeros(out_size,dtype=dtype),name='b')

    def _init_W(self):
        """根据输入输出大小初始化权重"""
        I,O = self.in_size,self.out_size
        W_data = np.random.randn(I,O).astype(self.dtype) * np.sqrt(1 / I)
        self.W.data = W_data

    def forward(self,x):
        #在传播数据时初始化权重
        if self.W.data is None:
            self.in_size = x.shape[1] #W的行数取决于X的列数
            self._init_W()
        y = F.linear(x,self.W,self.b)
        return y

#-------------------------------------
# 卷积层
#-------------------------------------
class Conv2d(Layer):
    def __init__(self, out_channels, kernel_size, stride=1,\
                 pad=0, nobias=False, dtype=np.float32, in_channels=None):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.pad = pad
        self.dtype = dtype

        #权重
        self.W = Parameter(None,name='W')
        if in_channels is not None:
            self._init_W()
        #偏差
        if nobias:
            self.b = None
        else:
            self.b = Parameter(np.zeros(out_channels, dtype=dtype),name='b')

    def _init_W(self, xp=np):
        C,OC = self.in_channels, self.out_channels
        KH,KW = dezero.utils.pair(self.kernel_size)
        #权值初始值使用Xavier初始值
        scale = np.sqrt( 1 / (C * KH * KW) )
        W_data = xp.random.randn(OC,C,KH,KW).astype(self.dtype) * scale
        self.W.data = W_data


    def forward(self,x):
        if self.W.data is None:
            self.in_channels = x.shape[1]
            xp = cuda.get_array_module(x)
            self._init_W(xp)

        y = F.conv2d_simple(x, self.W, self.b, self.stride, self.pad)
        #y = F.conv2d(x, self.W, self.b, self.stride, self.pad)
        return y


class RNN(Layer):
    """
    RNN 循环神经网络
    h := tanh(hW + xW +b)
    """
    def __init__(self, hidden_size, in_size=None):
        super().__init__()
        #x2h:将输入x转换为隐藏状态h的全连接层
        self.x2h = Linear(hidden_size,in_size=in_size)
        #h2h:将上一个隐藏状态转换为下一个隐藏状态的全连接层
        self.h2h = Linear(hidden_size,in_size=in_size, nobias=True)
        self.h = None #保存输出状态进行循环，注意第一轮没有输出状态

    def reset_state(self):
        self.h = None

    def forward(self,x):
        if self.h is None:
            h_new = F.tanh(self.x2h(x))
        else:
            h_new = F.tanh(self.x2h(x) + self.h2h(self.h))
        self.h = h_new
        return h_new
